import xml.dom.minidom as xdm
import csv
import openslide
import numpy as np
import os
from tqdm import trange
import cv2
from PIL import Image
import matplotlib.pyplot as plt
from time import time


# 读切片MSI类型标注的csv文件
def get_slide_type(filepath):
    with open(filepath) as f:
        reader = csv.reader(f)
        headers = next(reader)
        label_dict = {}
        for line in reader:
            label_dict[line[0]] = line[1]
    return label_dict


# 读肿瘤区域标注的xml文件
def get_tumor_edge(filename):
    domTree = xdm.parse(filename)
    collections = domTree.documentElement
    vertices_list = collections.getElementsByTagName('Vertices')
    multi_ver_list = []
    for k in range(len(vertices_list)):
        ver_list = []
        for i in range((len(vertices_list[k].childNodes) - 1) // 2):
            vertex = vertices_list[k].childNodes[2 * i + 1]  # 避开中间夹杂的换行文本数据
            x = int(vertex.getAttribute('X')) // (4**downsample_rate)
            y = int(vertex.getAttribute('Y')) // (4**downsample_rate)
            ver_list.append((x, y))
        multi_ver_list.append(ver_list.copy())
    return multi_ver_list


# 创建一个mask对非肿瘤区域进行涂黑，并生成mask后的图片
def get_masked_slide(source, labels):
    slide = openslide.OpenSlide(source)
    slide_size = slide.level_dimensions[downsample_rate]
    multi_ver_list = get_tumor_edge(labels)
    # 生成掩码mask
    mask = np.zeros((slide_size[1],slide_size[0],3), dtype=np.uint8)
    # 将顶点列表转换为opencv能处理的numpy数组
    for i, ver_list in enumerate(multi_ver_list):
        # opencv似乎不需要返回到初始顶点
        multi_ver_list[i] = np.array(ver_list[:-1])
    # 仅肿瘤区域为白色（255,255,255）
    cv2.fillPoly(mask, multi_ver_list, (255, 255, 255))
    slide_thumbnail = slide.get_thumbnail(slide_size)
    tile = np.array(slide_thumbnail)
    masked_slide = cv2.bitwise_and(tile, mask)

    return masked_slide


# 这里加入一个预处理操作，算出肿瘤区域的外包矩形顶点，切割时由此进行
def get_outer_rectangles(areas):  # 判断点，区域数组
    outer_rectangles = []
    for area in areas:
        min_x = min(area, key=lambda k: k[0])[0]
        min_y = min(area, key=lambda k: k[1])[1]
        max_x = max(area, key=lambda k: k[0])[0]
        max_y = max(area, key=lambda k: k[1])[1]
        outer_rectangles.append((min_x, min_y, max_x, max_y))
    return outer_rectangles


# 提取肿瘤区域的图块, 输入切片数据文件路径；顶点标注文件路径；切片标签；
# 图块大小；切割步进；肿瘤区域占比；保存图块间距；是否计算间距
def extract_tumor_patches(source, vertex_label, s_label,
                          patch_size=224, step=224, tumor_rate=0.9, distance=100, d_count=False):
    slide = get_masked_slide(source, vertex_label)
    slide_num = source[-6:-4]
    multi_ver_list = get_tumor_edge(vertex_label)
    outer_rects = get_outer_rectangles(multi_ver_list)
    # outer_rects = [[0, 0, slide.shape[1], slide.shape[0]]]  # 不计算外包矩形
    patch_num = 0
    savepoint = [0, 0]  # 记录保存图片的x y节点
    label_dir = os.path.join(s_label, slide_num)
    if not os.path.exists(label_dir):
        os.makedirs(label_dir)
    for outer_rect in outer_rects:
        for x in range(outer_rect[0], outer_rect[2]-patch_size, step):
            for y in range(outer_rect[1], outer_rect[3]-patch_size, step):
                patch = slide[y:y+patch_size, x:x+patch_size, :]  # 矩阵中x为列，y为行
                # 统计肿瘤区域像素点占比，即非（0,0,0）像素点占比
                if np.count_nonzero(np.any(patch, axis=2)) >= tumor_rate * patch_size * patch_size:
                    if d_count:
                        if (x - savepoint[0]) ** 2 + (y - savepoint[1]) ** 2 > distance ** 2:
                            savepoint = [x, y]
                            img = Image.fromarray(patch)
                            img.save(f"{label_dir}/tumor_patch_{slide_num}_{str(patch_num).zfill(5)}.png")
                            patch_num += 1
                    else:
                        img = Image.fromarray(patch)
                        img.save(f"{label_dir}/tumor_patch_{slide_num}_{str(patch_num).zfill(5)}.png")
                        patch_num += 1


if __name__ == '__main__':
    downsample_rate = 1  # 每级下采样4倍
    train_slide_num = 47  # 用于训练的切片数量
    start_num = 0
    source_path = "F:/PublicDatasets/PAIP2020/"
    slide_label_dict = get_slide_type(f"{source_path}labels/traning_data_MSI.csv")
    for i in trange(start_num, train_slide_num):
        file_code = f"training_data_{str(i+1).zfill(2)}"
        slide_path = f"{source_path}train/{file_code}.svs"
        vertices_path = f"{source_path}labels/training_data/{file_code}.xml"
        slide_label = slide_label_dict[file_code]
        extract_tumor_patches(slide_path, vertices_path, slide_label)

